# coding=UTF-8


import os
import numpy as np
import cv2
import glob
import matchfunc
from miscfunc import *
from python_sba import *
from triangulateMultiview import *
from viewSet import *
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plot

##################################
##相机内参
K=np.array([[1.037575214664696e+03, 0.00000000e+00,         6.422315830312182e+02],
            [0.00000000e+00,        1.043315752317925e+03,  3.878357750962377e+02],
            [0.00000000e+00,        0.00000000e+00,         1.0]])

RadialDistortion = [0.146911684283474,-0.214389634520344]
TangentialDistortion = [0.0, 0.0]

##################################
img_list=[]                 # 图像数据,M个
desc_list=[]                # 列表，每个图像的特征descriptor
kps_list=[]               # 列表，特征的keypoint
matches_dict={}         # (i-1,i)-->Matches, 前后两图像相互间的点匹配， M个图像产生M-1个匹配
goodMatches_dict={}     # (i-1,i)-->GoodMatches, 前后两图像相互间的点匹配，符合F矩阵的
relOrient_dict={}            # (i-1,i)-->Orient, 后一个图像的相机相对于前一个图像的相机的朝向     M*3*3
relLoc_dict={}               # (i-1,i)-->Loc, 后一个图像的相机相对于前一个图像的相机的位置     M*3*1

"""1，给定目录，读取图像文件名，编号为0,1,2,..."""
#files=glob.glob(r'D:\Program Files\MATLAB\R2016a\toolbox\vision\visiondata\structureFromMotion\*.jpg')
files=glob.glob(r'/home/wg/sfm_data/*.jpg')
M=len(files)     #总图像数量
#print(fileList)
####创建特征检测器
#surf = cv2.xfeatures2d.SURF_create(3000,nOctaves=8)  # 参数为Hessian阀值
sift = cv2.xfeatures2d.SIFT_create()


"""读取文件，特征检测"""
for file in files:
    img = cv2.imread(file)
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    ####畸变校正
    gray = cv2.undistort(gray, K, np.hstack((RadialDistortion, TangentialDistortion)))
    img_list.append(gray)
    kps, desc = sift.detectAndCompute(gray, None)
    kps_list.append(kps); desc_list.append(desc)        ##kps_list和desc_list从0到M对应各个视图的index
"""********建立视图集，加入第一个view"""
absOrient0=np.eye(3)
absLoc0=np.array([[0,0,0]]).transpose()
vSet=viewSet()              ##视图集
vSet.addView(0,absOrient0,absLoc0,kps_list[0],desc_list[0],(-1,[]))

"""********依次处理剩余的图像"""
M=len(files)
lastIdx=0
for idx in range(1,M):
    ####当前图像与前一个图像间进行特征点匹配
    matches = matchfunc.match_flann(desc_list[lastIdx], desc_list[idx])  # 进行flann匹配
    ##matches_dict[(i-1,i)]=matches
    #print("image%d,image%d:%d matches\n"%(idx-1,idx,len(matches)) )
    ##4, 计算当前图像与前一个图像间的F矩阵，内点
    N = len(matches)
    pts1, pts2 = matchfunc.fetchCoordinate(kps_list[lastIdx], kps_list[idx], matches)   #从keypoint提取坐标
    #   4.1 估计F矩阵
    F, mask = cv2.findFundamentalMat(pts1, pts2, cv2.FM_RANSAC,1, 0.99)
    #   4.2 得到符合F矩阵的内点
    goodPts1=np.array([pts1[i] for i in range(len(pts1)) if mask[i]])
    goodPts2=np.array([pts2[i] for i in range(len(pts2)) if mask[i]])
    goodMatches=[matches[i] for i in range(N) if mask[i]]
    #disp_matches(img_list[lastIdx],img_list[idx], kps_list[lastIdx], kps_list[idx], goodMatches)
    ##通过图像i-1与图像i的匹配点对相机i相对于相机i-1的方位（pose）进行估计
    R,t,ratio=matchfunc.estimateRelativePose(goodPts1,goodPts2,K,F)
    print("ratio=%f"%(ratio))
    if ratio<0.75:
        print("Pose of view%d can't be found! ratio=%f" % (idx,ratio))
        continue
    # 保存view i相对于view i-1的方位
    #####注意：R表示相机i-1坐标系到相机i坐标系的变换。因此，为了得到当前相机i相对于相机i-1的方位，我们必须做逆操作
    relOrient=R.transpose()
    relLoc=-np.dot(R.transpose(),t)
    #print(relLoc.transpose())
    #####计算当前相机相对于全局坐标系的方位（使用前一个相机的全局坐标系方位）
    lastAbsOrient,lastAbsLoc=vSet.getPoses(idx-1)    ##获取前一个view的绝对朝向和绝对位置
    curAbsOrient=np.dot(lastAbsOrient,relOrient)
    curAbsLoc=np.dot(lastAbsOrient,relLoc)+lastAbsLoc
    print(curAbsLoc.transpose())
    ##将当前view添加到视图集(只有正确的view才会加入)
    vSet.addView(idx,curAbsOrient,curAbsLoc,kps_list[idx],desc_list[idx],(lastIdx,goodMatches) )
    lastIdx=idx
    #####获取当前所有视图的方位，以及视图间的点轨迹
    camPoses=vSet.getAllPoses()
    tracks=vSet.generateTracks()
    #print(camPoses)
    #####使用多视图三角测量计算3D点
    pts3D=triangulateMultiview(tracks,camPoses,K)

    sbaCameras, sbaPoints=GenerateSBAdata(pts3D,tracks,camPoses,K)

    newcams, newpts, info = sba.SparseBundleAdjust(sbaCameras, sbaPoints)

    #print(np.sum(pts3D[:,2]>0)/len(pts3D))
    print(pts3D)

    fig =plot.figure()
    ax = fig.gca(projection='3d')
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("z")
    ax.plot(pts3D[:,0], pts3D[:,1], pts3D[:,2],'ko')
    ax.plot([0],[0],[0],'r.')
    plot.axis('equal')

    #####进行bundle adjustment
